package main

import (
	"image/color"

	"gonum.org/v1/plot"
	"gonum.org/v1/plot/plotter"
	"gonum.org/v1/plot/vg"
)

func main() {
	p := plot.New()
	p.Title.Text = "CUDA Warp Shuffle Data Flow"
	p.X.Label.Text = "Thread ID"
	p.Y.Label.Text = "Value Source Thread ID"

	// 画 shfl_up_sync(delta=1)
	upPoints := make(plotter.XYs, 0)
	for tid := 1; tid < 32; tid++ {
		upPoints = append(upPoints, plotter.XY{X: float64(tid), Y: float64(tid)})
		upPoints = append(upPoints, plotter.XY{X: float64(tid), Y: float64(tid - 1)})
	}
	upLine, _ := plotter.NewLine(upPoints)
	upLine.Color = color.RGBA{B: 255, A: 255} // blue
	upLine.Width = vg.Points(1)
	p.Add(upLine)

	// 画 shfl_down_sync(delta=1)
	downPoints := make(plotter.XYs, 0)
	for tid := 0; tid < 31; tid++ {
		downPoints = append(downPoints, plotter.XY{X: float64(tid), Y: float64(tid)})
		downPoints = append(downPoints, plotter.XY{X: float64(tid), Y: float64(tid + 1)})
	}
	downLine, _ := plotter.NewLine(downPoints)
	downLine.Color = color.RGBA{R: 255, A: 255} // red
	downLine.Width = vg.Points(1)
	p.Add(downLine)

	// 添加图例
	p.Legend.Add("shfl_up_sync(delta=1)", upLine)
	p.Legend.Add("shfl_down_sync(delta=1)", downLine)
	p.Legend.Top = true

	// 保存为 PNG
	if err := p.Save(8*vg.Inch, 4*vg.Inch, "shuffle_flow.png"); err != nil {
		panic(err)
	}
}
